import { app } from "/scripts/app.js"
import { api } from "/scripts/api.js"

(function() {

// Workflow publishing
// - done whenever the graph changes, as long as there is a KritaOutput node

let publisherRegistered = false

async function publishWorkflow(e) {
    const prompt = await app.graphToPrompt()
    await api.fetchApi("/api/etn/workflow/publish", {
        method: "POST",
        body: JSON.stringify({
            name: "ComfyUI Web",
            client_id: api.clientId,
            workflow: prompt["output"]
        }, null, 2)
    })
}


// Image background for nodes
// - this is just for visuals

function loadImage(base64) {
    const image = new Image()
    image.src = base64
    // image.onerror = () => console.error("Failed to load image");
    return image
}

const canvasIcon = loadImage("")
const outputIcon = loadImage("")

function setIconImage(nodeType, image, size, padRows, padCols) {
    const onAdded = nodeType.prototype.onAdded
    nodeType.prototype.onAdded = function () {
        onAdded?.apply(this, arguments)
        this.size = size
    }

    const onDrawBackground = nodeType.prototype.onDrawBackground
    nodeType.prototype.onDrawBackground = function(ctx) {
        onDrawBackground?.apply(this, arguments)

        const pad = [padCols * 20, LiteGraph.NODE_SLOT_HEIGHT * padRows + 8];
        if(this.flags.collapsed || pad[1] + 32 > this.size[1] || image.width === 0) {
            return
        }
        const avail = [this.size[0] - pad[0], this.size[1] - pad[1]]
        const scale = Math.min(1.0, avail[0] / image.width, avail[1] / image.height)
        const size = [Math.floor(image.width * scale), Math.floor(image.height * scale)]
        const offset = [Math.max(0, (avail[0] - size[0]) / 2), Math.max(0, (avail[1] - size[1]) / 2)]
        ctx.drawImage(image, offset[0], pad[1] + offset[1], size[0], size[1])
    }
}


// Parameter node
// - represents a customizable parameter that should be exposed in external tools
// - adapts to whichever node it is connected to, similar to the built-in "Primitive" node
// - can only be connected to slots which are converted widgets

const replaceableWidgets = ["INT", "FLOAT", "BOOLEAN", "STRING", "COMBO", "INT:seed"]

const parameterTypes = {
    "combo": ["choice"],
    "number": ["number", "number (integer)"],
    "toggle": ["toggle"],
    "text": ["text", "prompt (positive)", "prompt (negative)"],
}

function defaultParameterType(widgetType, connectedNode, connectedWidget) {
    let paramType = parameterTypes[widgetType][0]
    if (connectedNode.comfyClass === "CLIPTextEncode") {
        paramType = "prompt (positive)"
    }
    if (connectedWidget.options?.round === 1) {
        paramType = "number (integer)"
    }
    return paramType
}

function valueMatchesType(value, type, options) {
    if (type === "number") {
        return typeof value === "number"
    } else if (type === "combo") {
        return options?.values?.includes(value)
    } else if (type === "toggle") {
        return typeof value === "boolean"
    }
    return typeof value === "string"
}

function optionalWidgetValue(widgets, index, fallback) {
    const result = widgets.length > index ? widgets[index].value : null
    return result === null || result === 0 ? fallback : result
}

function changeWidgets(node, type, connectedNode, connectedWidget) {
    if (type === "customtext") {
        type = "text"
    }
    const options = connectedWidget.options

    const parameterTypeHint = node.widgets[1].value
    const notSpecialized = node.widgets[1].options.values.includes("auto")
    const parameterTypeMismatch = !parameterTypes[type].includes(parameterTypeHint)
    if (notSpecialized || parameterTypeMismatch) {
        node.widgets[1].options = {values: parameterTypes[type]}
    }
    if (parameterTypeMismatch) {
        node.widgets[1].value = defaultParameterType(type, connectedNode, connectedWidget)
    }
    const oldDefault = node.widgets.length > 2 ? node.widgets[2].value : connectedWidget.value
    const oldMin = optionalWidgetValue(node.widgets, 3, options?.min ?? 0)
    const oldMax = optionalWidgetValue(node.widgets, 4, options?.max ?? 100)
    const isDefaultValid = valueMatchesType(oldDefault, type, connectedWidget.options)
    while (node.widgets.length > 2) {
        node.widgets.pop()
    }
    const value = isDefaultValid && oldDefault !== "" ? oldDefault : connectedWidget.value
    node.addWidget(type, "default", value, null, options)
    if (type === "number") {
        node.addWidget("number", "min", oldMin, null, options)
        node.addWidget("number", "max", oldMax, null, options)
    }
}

function adaptWidgetsToConnection(node) {
    if (!node.outputs || node.outputs.length === 0) {
        return
    }
    const links = node.outputs[0].links
    if (links && links.length === 1) {
        const link = node.graph.links[links[0]]
        if (!link) return

        const theirNode = node.graph.getNodeById(link.target_id)
        if (!theirNode || !theirNode.inputs) return
    
        const input = theirNode.inputs[link.target_slot]
        if (!input || theirNode.widgets === undefined) return

        node.outputs[0].type = input.type

        if (node.widgets[0].value === "Parameter") {
            node.widgets[0].value = input.name
        }

        const widgetName = input.widget.name
        const theirWidget = theirNode.widgets.find((w) => w.name === widgetName)
        if (!theirWidget) return // connected to a custom node that isn't installed

        const widgetType = theirWidget.origType ?? theirWidget.type
        changeWidgets(node, widgetType, theirNode, theirWidget)

    } else if (!links || links.length === 0) {
        node.outputs[0].type = "*"
        node.widgets[1].value = "auto"
        node.widgets[1].options = {values: ["auto"]}
    }
}

function setupParameterNode(nodeType) {
    const onAdded = nodeType.prototype.onAdded
    nodeType.prototype.onAdded = function() {
        onAdded?.apply(this, arguments)
        adaptWidgetsToConnection(this)
    }

    const onAfterGraphConfigured = nodeType.prototype.onAfterGraphConfigured
    nodeType.prototype.onAfterGraphConfigured = function() {
        onAfterGraphConfigured?.apply(this, arguments)
        adaptWidgetsToConnection(this)
    }

    const onConnectOutput = nodeType.prototype.onConnectOutput
    nodeType.prototype.onConnectOutput = function(slot, type, input, target_node, target_slot) {
        if (!input.widget && !(input.type in replaceableWidgets)) {
            return false
        } else if (onConnectOutput) {
            result = onConnectOutput.apply(this, arguments)
            return result
        }
        return true
    }

    const onConnectionsChange = nodeType.prototype.onConnectionsChange
    nodeType.prototype.onConnectionsChange = function(_, index, connected) {
        if (!app.configuringGraph) {
            adaptWidgetsToConnection(this)
        }
        onConnectionsChange?.apply(this, arguments)
    }
}


// Register the extension

app.registerExtension({
	name: "external_tooling_nodes",

    beforeRegisterNodeDef(nodeType /*typeof LGraphNode*/, nodeData /*ComfyObjectInfo*/, app) {
        if (nodeData.name === "ETN_KritaCanvas") {
            setIconImage(nodeType, canvasIcon, [200, 100], 0, 2)
        } else if (nodeData.name === "ETN_KritaOutput") {
            setIconImage(nodeType, outputIcon, [200, 100], 1, 0)
        } else if (nodeData.name === "ETN_Parameter") {
            setupParameterNode(nodeType)
        } else if (nodeData.name === "ETN_SendText") {
            const onAdded = nodeType.prototype.onAdded
            nodeType.prototype.onAdded = function() {
                onAdded?.apply(this, arguments)
                this.inputs[0].type = "*"
            }
        }
    },

    nodeCreated(node /*ComfyNode*/, app) {
        if (publisherRegistered || node.comfyClass !== "ETN_KritaOutput") {
            return
        }
        api.addEventListener('graphChanged', publishWorkflow)
        publisherRegistered = true
    },

    setup(app) {
        if (publisherRegistered) {
            publishWorkflow(null)
        }
    },
});

})();